Import Necessary Modules

from pathlib import Path

from ipyniivue import download_dataset

BASE_API_URL = "https://niivue.com/demos/images/"
DATA_FOLDER = Path("images")

# Download data for example
download_dataset(
    BASE_API_URL,
    dest_folder=DATA_FOLDER,
    files=[
        "mni152.nii.gz",
        "mni152_pveseg.nii.gz",
    ],
)

mni152.nii.gz already exists.
Downloading mni152_pveseg.nii.gz...
Dataset downloaded successfully to images.

import json

import ipywidgets as widgets

from ipyniivue import NiiVue, ShowRender, SliceType

# Setup NiiVue instance

nv = NiiVue(
    back_color=(0.3, 0.3, 0.3, 1),
    show_3d_crosshair=True,
)
nv.set_radiological_convention(False)

nv.load_volumes(
    [
        {"path": DATA_FOLDER / "mni152.nii.gz"},
        {"path": DATA_FOLDER / "mni152_pveseg.nii.gz", "opacity": 0.5},
    ]
)

nv.opts.multiplanar_show_render = ShowRender.ALWAYS
nv.set_slice_type(SliceType.MULTIPLANAR)
nv.graph.auto_size_multiplanar = True

# Add widgets

background_checkbox = widgets.Checkbox(
    value=True,
    description="Background",
)

mask_checkbox = widgets.Checkbox(
    value=False,
    description="Mask",
)

smooth_checkbox = widgets.Checkbox(
    value=True,
    description="Smooth",
)

opacity_slider = widgets.IntSlider(
    value=127,
    min=1,
    max=255,
    description="Opacity",
)

# Define Event Handlers


def on_opacity_change(change):
    """Adjust the opacity of the second volume."""
    nv.set_opacity(1, change["new"] / 255)


def on_background_checkbox_change(change):
    """Show or hide the background volume."""
    nv.volumes[0].opacity = 1.0 if change["new"] else 0.0


def on_mask_checkbox_change(change):
    """Set is_alpha_clip_dark."""
    nv.opts.is_alpha_clip_dark = change["new"]


def on_smooth_checkbox_change(change):
    """Set interpolation."""
    nv.set_interpolation(not change["new"])


opacity_slider.observe(on_opacity_change, names="value")
background_checkbox.observe(on_background_checkbox_change, names="value")
mask_checkbox.observe(on_mask_checkbox_change, names="value")
smooth_checkbox.observe(on_smooth_checkbox_change, names="value")

# Set Initial Values

on_opacity_change({"new": opacity_slider.value})
on_background_checkbox_change({"new": background_checkbox.value})
on_mask_checkbox_change({"new": mask_checkbox.value})
on_smooth_checkbox_change({"new": smooth_checkbox.value})

# Handle location changes

location_label = widgets.HTML(" ")


def handle_location_change(data):
    """Update the location label when the crosshair moves."""
    location_label.value = "  " + data["string"]


nv.on_location_change(handle_location_change)

# Handle Custom Colormap Input

colormap_textarea = widgets.Textarea(
    value="""{
    "R": [0,   0, 120, 175],
    "G": [0,  90,  60, 185],
    "B": [0, 120,  60, 175],
    "labels": ["air", "CSF", "gray", "white"]
}""",
    description="Custom Colormap (JSON)",
    layout=widgets.Layout(width="100%", height="150px"),
)

apply_button = widgets.Button(
    description="Apply Colormap",
)


def on_apply_button_click(b):
    """Apply custom colormap json."""
    json_str = colormap_textarea.value
    try:
        cmap = json.loads(json_str)
        nv.volumes[1].set_colormap_label(cmap)
    except json.JSONDecodeError as e:
        print("Invalid JSON format:", e)


apply_button.on_click(on_apply_button_click)

on_apply_button_click(None)

# Display all

controls = widgets.HBox(
    [
        background_checkbox,
        mask_checkbox,
        smooth_checkbox,
        opacity_slider,
    ]
)

script_controls = widgets.VBox(
    [
        colormap_textarea,
        apply_button,
    ]
)

display_widgets = widgets.VBox(
    [
        controls,
        nv,
        location_label,
        script_controls,
    ]
)

display(display_widgets)